import tensorflow as tf
import numpy as np
import gzip


def read_idx3(filename):
    """
    读取gz格式的数据集图像部分，并返回

    :param filename: extension name of the file is '.gz'
    :return: images data, shape -> num, rows, cols
    """
    with gzip.open(filename) as fo:
        print('Reading images...')
        buf = fo.read()

        offset = 0  # 偏移量
        # 首先获取的是这个数据集的头部数据，通常是元数据。
        #   '>i'  表示顺序读取，并且数据类型为整数
        #   4  读4个单位
        #   offset 偏移量
        # 返回的是一个数组，赋值给header
        header = np.frombuffer(buf, dtype='>i', count=4, offset=offset)
        print(header)
        magic_number, num_images, num_rows, num_cols = header
        # magic number 即幻数，意义不明，只是读取时需要占位所以声明了
        print("\tmagic number: {}, number of images: {}, number of rows: {}, number of columns: {}" \
              .format(magic_number, num_images, num_rows, num_cols))
        # 计算偏移量，以读取后续的内容
        # size = 数组长度
        # itemsize = 每个元素的大小
        # 因此乘起来就是跳过header的内容，读后续的内容
        offset += header.size * header.itemsize
        # 读取真正的数据。>B 表示是二进制数据
        data = np.frombuffer(buf, '>B', num_images * num_rows * num_cols, offset).reshape(
            (num_images, num_rows, num_cols))
        # .reshape 表示按传入的参数重新构造这个数组

        return data, num_images


def read_idx1(filename):
    """
    读取gz格式的数据集标签部分，并返回

    :param filename: extension name of the file is '.gz'
    :return: labels
    """
    with gzip.open(filename) as fo:
        print('Reading labels...')
        buf = fo.read()

        offset = 0
        header = np.frombuffer(buf, '>i', 2, offset)
        magic_number, num_labels = header
        print("\tmagic number: {}, number of labels: {}" \
              .format(magic_number, num_labels))

        offset += header.size * header.itemsize

        data = np.frombuffer(buf, '>B', num_labels, offset)
        return data


data_source_path = 'D:/emnist/'
train_labels = read_idx1(data_source_path + 'emnist-letters-train-labels-idx1-ubyte.gz')
train_images, train_images_num = read_idx3(data_source_path + 'emnist-letters-train-images-idx3-ubyte.gz')
test_labels = read_idx1(data_source_path + 'emnist-letters-test-labels-idx1-ubyte.gz')
test_images, test_images_num = read_idx3(data_source_path + 'emnist-letters-test-images-idx3-ubyte.gz')

# reshape()中的第一个参数传入训练集（测试集）的大小，下同。
train_images = train_images.reshape((train_images_num, 28, 28, 1))
test_images = test_images.reshape((test_images_num, 28, 28, 1))

# 像素值映射到 0 - 1 之间
train_images, test_images = train_images / 255.0, test_images / 255.0
# 标签值处理：原始标签值范围1~26，改为0~25
train_labels = train_labels - 1
test_labels = test_labels - 1

model = tf.keras.models.Sequential([
    # 第1个卷积层
    tf.keras.layers.Conv2D(filters=32,
                           kernel_size=(3, 3),
                           input_shape=(28, 28, 1),
                           activation='relu',
                           padding='same'),
    # 防止过拟合
    # tf.keras.layers.Dropout(rate=0.3),
    # 第1个池化层
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    # 第2个卷积层
    tf.keras.layers.Conv2D(filters=64,
                           kernel_size=(3, 3),
                           activation='relu',
                           padding='same'),
    # 防止过拟合
    # tf.keras.layers.Dropout(rate=0.3),
    # 第2个池化层
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    # 平坦层
    tf.keras.layers.Flatten(),
    # 输出层
    tf.keras.layers.Dense(26, activation='softmax')
])

# ================= 模型训练 ===================
# 设置训练参数
training_epochs = 5  # 训练轮数
batch_size = 100  # 单次训练样本数
# 定义训练模式
model.compile(optimizer='adam',  # 优化器
              loss=tf.keras.losses.sparse_categorical_crossentropy,  # 损失函数
              metrics=['accuracy'])  # 评估指标
# 模型训练
train_history = model.fit(x=train_images, y=train_labels,
          validation_split=0.2,
          epochs=training_epochs,
          batch_size=batch_size,
          verbose=2)


# ================= 评估模型 ===================
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print('评估模型：')
print(test_loss, test_acc)

# ================= 应用模型 ===================
print('应用模型')
preds = model.predict_classes(test_images)
print(preds)
